/**
 * 
 */
package edu.umd.clip.lm.programs;

import java.io.*;
import java.nio.channels.*;

import javax.xml.parsers.*;

import org.w3c.dom.*;
import org.xml.sax.SAXException;

import edu.berkeley.nlp.util.*;
import edu.umd.clip.jobs.JobManager;
import edu.umd.clip.lm.model.*;
import edu.umd.clip.lm.model.data.OnDiskTrainingDataReader;
import edu.umd.clip.lm.model.data.ReadableTrainingData;
import edu.umd.clip.lm.model.data.TrainingDataReader;
import edu.umd.clip.lm.model.training.*;
import edu.umd.clip.lm.util.XMLUtils;
import edu.umd.clip.lm.util.tree.BinaryTree;

/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class LMNewTrainer {

	public static class Options {
        @Option(name = "-config", required = true, usage = "XML config file")
		public String config;
        @Option(name = "-jobs", usage = "number of concurrent jobs (default: 1)")
        public int jobs = 1;
        @Option(name = "-lm", required = false, usage = "LM ID to train (default: " + LanguageModel.PRIMARY_LM_ID + ")")
		public String lm = LanguageModel.PRIMARY_LM_ID;      
        @Option(name = "-load", usage = "load a snapshot")
		public String load;
        @Option(name = "-tmpdir", usage = "directory for temporary files (default: /tmp)")
		public String tmpDir = "/tmp";
        @Option(name = "-nr-exchange-init", usage = "number of random initializations of the Exchange algorithm (default: 1)")
		public int nrExchangeInits = 1;
	}

	/**
	 * @param args
	 * @throws IOException 
	 */
	public static void main(String[] args) throws IOException {
        OptionParser optParser = new OptionParser(Options.class);
        Options opts = (Options) optParser.parse(args, true);

		Experiment.initialize(opts.config);
		Experiment experiment = Experiment.getInstance();
		LanguageModel lm = experiment.getLM(opts.lm);

		JobManager.initialize(opts.jobs);
		Thread thread = new Thread(JobManager.getInstance(), "Job Manager");
		thread.setDaemon(true);
		thread.start();
		
		experiment.buildPrefixes();
		
		File tmpDir = new File(opts.tmpDir);
		if (!tmpDir.isDirectory()) {
			tmpDir.mkdirs();
		}
		NewTrainer.setNR_RANDOM_EXCHANGE_INIT(opts.nrExchangeInits);
		NewTrainer trainer = new NewTrainer(lm, tmpDir);

		TrainingContext context;
		context = trainer.initialize();
		
		if (opts.load != null) {
			System.out.printf("loading snapshot %s\n", opts.load);
			DocumentBuilderFactory dbf = DocumentBuilderFactory.newInstance();
			DocumentBuilder db = null;
			try {
				db = dbf.newDocumentBuilder();
				Document doc = db.parse(opts.load);
				
				BinaryTree<HistoryTreePayload> snapshot = new BinaryTree<HistoryTreePayload>(doc.getDocumentElement());
				doc = null;
				db = null;
				context = trainer.loadSnapshot(context, snapshot);
			}catch(ParserConfigurationException pce) {
				pce.printStackTrace();
				System.exit(1);
			}catch(SAXException se) {
				se.printStackTrace();
			}catch(IOException ioe) {
				ioe.printStackTrace();
			}

		}

		trainer.train(context);
		
	}

}
